# Set the working directory to a specific path
output_dir <- "yourpath/"
dir.create(output_dir, recursive = TRUE, showWarnings = FALSE)
setwd(output_dir)

# Ensure necessary libraries are installed and loaded
if (!require("minpack.lm")) install.packages("minpack.lm")
if (!require("ggplot2")) install.packages("ggplot2")
if (!require("readxl")) install.packages("readxl")
if (!require("writexl")) install.packages("writexl")

library(minpack.lm)
library(ggplot2)
library(readxl)
library(writexl)

# Load data from the Excel file from a specific sheet
data_path <- "yourpath/070724_IDL_equations_results_rearrange_48.xlsx"
data <- read_excel(data_path, sheet = "Sheet2")

# Prepare the results data frame
results <- data.frame(State = character(), k = integer(), MSE = numeric(), Lambda = numeric(), DelayLength = numeric(), stringsAsFactors = FALSE)

# Iterate over each state's data
for (i in 1:nrow(data)) {
  state_name <- data$State[i]
  coeffs <- as.numeric(data[i, 2:10])  # Assume the coefficients are in columns 2 to 10
  x <- 0:8  # Use the first 9 lags
  
  # Fit models and save results
  fit_models <- list()
  mse_results <- c()
  lambda_values <- c()
  delay_lengths <- c()
  plots <- list()
  
  for (k in 1:3) {
    formula <- as.formula(sprintf("y ~ lambda^%d * x^%d * exp(-lambda * x) / factorial(%d)", k, k-1, k-1))
    fit <- nls(formula, start = list(lambda = 1), data = data.frame(x = x, y = coeffs), 
               algorithm = "port", control = nls.control(maxiter = 1000))
    fit_models[[k]] <- fit
    mse <- mean((coeffs - predict(fit))^2)
    lambda <- coef(fit)["lambda"]
    mse_results <- c(mse_results, mse)
    lambda_values <- c(lambda_values, lambda)
    delay_length <- k / lambda
    delay_lengths <- c(delay_lengths, delay_length)
    
    # Create a higher density x sequence for plotting
    x_seq <- seq(min(x), max(x), length.out = 1000)
    
    # Plot
    plots[[k]] <- ggplot(data = data.frame(x = x, y = coeffs), aes(x, y)) +
      geom_point(size = 4, shape = 16) +
      geom_line(data = data.frame(x = x_seq, y = predict(fit, newdata = data.frame(x = x_seq))), color = 'blue', size = 1.5) +
      xlab("Lag") + ylab("Coefficient") +
      ggtitle(state_name) +  # Add state name as the title
      theme_classic(base_size = 20) +  # White background with only axis lines and larger font size
      theme(
        plot.title = element_text(size = 30, face = "bold", hjust = 0.5, vjust = 0.5, margin = margin(b = 20)),  # Centered and bold title with larger font size
        axis.title.x = element_text(size = 28),  # Increase font size for x-axis label
        axis.title.y = element_text(size = 28),  # Increase font size for y-axis label
        axis.text.x = element_text(size = 28, color = "black"),  # Increase font size for x-axis text
        axis.text.y = element_text(size = 28, color = "black")   # Increase font size for y-axis text
      ) +
      annotate("text", x = max(x) - 1, y = max(coeffs) * 0.88,  # Adjust the position to be slightly lower
               label = sprintf("Delay Order k = %d\nDelay Length = %.4f", k, delay_length), hjust = 1, size = 9)  # Adjusted vertical position of annotation
  }
  
  # Find the model with the minimum MSE
  min_mse_index <- which.min(mse_results)
  min_mse <- mse_results[min_mse_index]
  min_lambda <- lambda_values[min_mse_index]
  min_delay_length <- delay_lengths[min_mse_index]
  
  # Save the plot with the minimum MSE
  ggsave(sprintf("%s_plot.png", state_name, min_mse_index), plot = plots[[min_mse_index]], width = 10, height = 8)
  
  # Store the results in the dataframe
  results <- rbind(results, data.frame(State = state_name, k = min_mse_index, MSE = min_mse, Lambda = min_lambda, DelayLength = min_delay_length))
}
